package main

import (
	"encoding/json"
	"log"
	"time"

	"github.com/shirou/gopsutil/disk"
)

// 磁盘 I/O 信息结构体（与你定义一致）
type DiskIocInfo struct {
	ReadCount        uint64 `json:"readCount"`
	MergedReadCount  uint64 `json:"mergedReadCount"`
	WriteCount       uint64 `json:"writeCount"`
	MergedWriteCount uint64 `json:"mergedWriteCount"`
	ReadBytes        uint64 `json:"readBytes"`  // 累计读取字节数
	WriteBytes       uint64 `json:"writeBytes"` // 累计写入字节数
	ReadTime         uint64 `json:"readTime"`
	WriteTime        uint64 `json:"writeTime"`
	IoTime           uint64 `json:"ioTime"`
	WeightedIO       uint64 `json:"weightedIO"`
	Name             string `json:"name"` // 磁盘名称（如 sda）
	Label            string `json:"label"`
}

func (u DiskIocInfo) String() string {
	s, _ := json.Marshal(u)
	return string(s)
}

// 存储上一次的磁盘 I/O 统计数据（key: 磁盘名称，value: 对应 I/O 信息）
var lastDiskIOC map[string]DiskIocInfo

func main() {
	// 初始化上一次数据存储
	lastDiskIOC = make(map[string]DiskIocInfo)
	// 采集间隔（可调整，比如 1 秒）
	interval := 2000 * time.Millisecond

	for {
		// 1. 获取当前磁盘 I/O 统计
		diskiocs, err := disk.IOCounters()
		if err != nil {
			log.Printf("获取磁盘 I/O 统计失败：%v", err)
			time.Sleep(interval)
			continue
		}

		// 2. 遍历当前数据，计算每个磁盘的读写速率
		for diskName, io := range diskiocs {
			// 构造当前磁盘的 I/O 信息
			currentIO := DiskIocInfo{
				ReadCount:        io.ReadCount,
				MergedReadCount:  io.MergedReadCount,
				WriteCount:       io.WriteCount,
				MergedWriteCount: io.MergedWriteCount,
				ReadBytes:        io.ReadBytes,
				WriteBytes:       io.WriteBytes,
				ReadTime:         io.ReadTime,
				WriteTime:        io.WriteTime,
				IoTime:           io.IoTime,
				WeightedIO:       io.WeightedIO,
				Name:             diskName,
				Label:            io.Label,
			}

			// 3. 计算速率：如果有上一次数据，则计算差值/时间间隔
			if lastIO, exists := lastDiskIOC[diskName]; exists {
				// 计算读写字节数差值
				log.Println(currentIO.ReadBytes, lastIO.ReadBytes)
				readDiff := currentIO.ReadBytes - lastIO.ReadBytes
				writeDiff := currentIO.WriteBytes - lastIO.WriteBytes

				// 计算每秒速率（字节/秒）
				readSpeed := float64(readDiff) / interval.Seconds()
				writeSpeed := float64(writeDiff) / interval.Seconds()

				// 打印结果（格式：磁盘名称 | 读取速率 | 写入速率）
				log.Printf(
					"磁盘: %s | 读取速率: %.2f B/s (%.2f KB/s) | 写入速率: %.2f B/s (%.2f KB/s)",
					diskName,
					readSpeed, readSpeed/1024, // 转换为 KB/s 更直观
					writeSpeed, writeSpeed/1024,
				)
			} else {
				// 第一次采集，无历史数据，仅提示
				log.Printf("磁盘: %s | 首次采集，暂无速率数据", diskName)
			}

			// 4. 更新上一次数据：将当前数据存入 lastDiskIOC
			lastDiskIOC[diskName] = currentIO
		}

		// 等待下一次采集
		time.Sleep(interval)
	}
}
